package com.lizba.redis.limit.funnel;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * <p>
 *      漏斗限流
 * </p>
 *
 * @Author: Liziba
 * @Date: 2021/9/7 20:24
 */
public class FunnelRateLimiter {

    /** map用于存储多个漏斗 */
    private Map<String, Funnel> funnels = new ConcurrentHashMap<>();



    /**
     * 请求（行为）是否被允许
     *
     * @param userId        用户id
     * @param actionKey     行为key
     * @param capacity      漏斗容量
     * @param leakingRate   剩余容量
     * @param quota         请求次数
     * @return
     */
    public boolean isActionAllowed(String userId, String actionKey, int capacity, float leakingRate, int quota) {

        String key = String.format("%s:%s", userId, actionKey);
        Funnel funnel = funnels.get(key);
        if (funnel == null) {
            funnel = new Funnel(capacity, leakingRate);
            funnels.put(key, funnel);
        }
        return funnel.waterLeaking(quota);
    }


    /**
     * 漏斗类
     */
    class Funnel {
        /** 漏斗容量 */
        int capacity;
        /** 漏斗流速，每毫秒允许的流速（请求） */
        float leakingRate;
        /** 漏斗剩余空间 */
        int leftCapacity;
        /** 上次漏水时间 */
        long leakingTs;

        public Funnel(int capacity, float leakingRate) {
            this.capacity = this.leftCapacity = capacity;
            this.leakingRate = leakingRate;
            leakingTs = System.currentTimeMillis();
        }

        /**
         * 计算剩余空间
         */
        void makeSpace() {
            long nowTs = System.currentTimeMillis();
            long intervalTs = nowTs - leakingTs;
            int intervalCapacity = (int) (intervalTs * leakingRate);
            // int 溢出
            if (intervalCapacity < 0) {
                this.leftCapacity = this.capacity;
                this.leakingTs = nowTs;
                return;
            }
            // 腾出空间 >= 1
            if (intervalCapacity < 1) {
                return;
            }
            // 增加漏斗剩余容量
            this.leftCapacity += intervalCapacity;
            this.leakingTs = nowTs;
            // 容量不允许超出漏斗容量
            if (this.leftCapacity > this.capacity) {
                this.leftCapacity = this.capacity;
            }
        }

        /**
         * 漏斗流水
         *
         * @param quota     流水量
         * @return
         */
        boolean waterLeaking(int quota) {
            // 触发漏斗流水
            this.makeSpace();
            if (this.leftCapacity >= quota) {
                leftCapacity -= quota;
                return true;
            }
            System.out.println(leftCapacity);
            return false;
        }
        
    }

}
